# -*- coding: utf-8 -*-
"""
Created on 2019/11/22 上午11:20

@Project -> File: ode-neural-network -> ode_net.py

@Author: luolei

@Describe: 训练ODE Net模型
"""

import matplotlib.pyplot as plt
import torch
import pandas as pd
import numpy as np
import copy
import json

from lib import discrete_t_steps, cols, lr, epochs, cols_n
from lib import normalize
from mod.neural_network_model.nn import PartialDeriveNet
from mod.neural_network_model.loss_criterion import criterion


def get_integrate_t_arr(t0, dt, discrete_t_steps):
	"""
	获取各个点的积分时刻表
	:param t0: torch.tensor, 起始时间序列
	:param dt: torch.tensor, 时间间隔
	:param discrete_t_steps: int, 起始时间和结束时间离散个数
	:return: integrate_t_arr, torch.tensor, 积分各时刻表
	"""
	
	t, integrate_t_arr = copy.deepcopy(t0), copy.deepcopy(t0)
	for step in range(discrete_t_steps):
		t += dt
		integrate_t_arr = np.hstack((integrate_t_arr, t))
	return integrate_t_arr


def ODESolver(x0, integrate_t_arr, dt_tensor, pd_net):
	"""
	ODE积分求解器
	:param x0: torch.tensor, shape = (points, dim_x)
	:param integrate_t_arr: torch.tensor, 用于确定每次积分步起止时刻的时间表，shape = (points, discrete_t_steps)
	:param dt_tensor: torch.tensor, shape = (points, dim_x)
	:param pd_net: torch.nn object, partial derive network obtained
	:return: x: torch.tensor, final integrated x, shape = (points, dim_x)
	:return: x_records: torch.tensor, records of x, shape = (time, points, dim_x)
	"""
	
	x = copy.deepcopy(x0)  # 初值
	x_records = x.unsqueeze(0)
	for step in range(integrate_t_arr.shape[1] - 1):  # **注意积分计算次数
		input = torch.cat((x, integrate_t_arr[:, step: step + 1]), dim = 1)
		partial_derives = pd_net(input)
		
		delta_x = torch.mul(partial_derives, dt_tensor)  # **点乘
		x = torch.add(x, delta_x)
		x_records = torch.cat((x_records, x.unsqueeze(0)), dim = 0)
	
	return x, x_records


def integrate(x0, t0, t1, dt, pd_net):
	"""
	计算初值问题的解
	:param x0: torch.tensor, shape = (points, dim_x)
	:param t0: float, initial time
	:param t1: float, final time
	:param dt: float, time interval len for integration
	:param pd_net: torch.nn object, partial derive net obtained
	:return: x: torch.tensor, final integrated x, shape = (points, dim_x)
	:return: x_records: torch.tensor, records of x, shape = (time, points, dim_x)
	"""
	
	integrate_t_arr = np.arange(t0, t1, dt).reshape(1, -1)
	integrate_t_arr = torch.from_numpy(integrate_t_arr.astype(np.float32))
	
	dt = np.array([dt]).reshape(-1, 1)
	dt = torch.from_numpy(dt.astype(np.float32))
	dt_tensor = dt.mm(torch.ones(1, 2))
	
	x, x_records = ODESolver(x0, integrate_t_arr, dt_tensor, pd_net)
	
	return x, x_records


def save_models(pd_net, train_loss_record):
	"""保存模型文件"""
	
	# 保存模型文件
	torch.save(pd_net.state_dict(), '../../file/model/state_dict.pth')
	
	# 保存模型结构参数
	model_struc_params = {
		'pd_net': {
			'input_size': pd_net.input_size,
			'hidden_sizes': pd_net.hidden_sizes,
			'output_size': pd_net.output_size
		}
	}
	
	with open('../../file/model/struc_params.json', 'w') as f:
		json.dump(model_struc_params, f)
		
	# 保存损失函数记录
	train_loss_list = [float(p.detach().cpu().numpy()) for p in train_loss_record]
	
	with open('../../file/model/train_loss.json', 'w') as f:
		json.dump(train_loss_list, f)
		
		
def vstack_arr(total, sub):
	"""竖向拼接数据"""
	
	if total is None:
		total = sub
	else:
		total = np.vstack((total, sub))
		
	return total
	

if __name__ == '__main__':
	# %% 载入数据
	data = pd.read_csv('../../data/runtime/total_obs_data.csv')

	# 数据归一化
	data = normalize(data)
	
	# %% 构造数据样本
	labels = list(data['label'].drop_duplicates())
	x0, x1 = None, None
	dt = None
	integrate_t_arr = None
	for label in labels:
		sub_data = data[data.label == label].copy()
		
		sub_x0 = sub_data.iloc[: -1][cols].to_numpy()   # 获得初值x0
		sub_x1 = sub_data.iloc[1:][cols].to_numpy()     # 获得终值x1
		
		sub_t0 = sub_data.iloc[: -1][['time']].to_numpy()   # 获得初始时间t0
		sub_t1 = sub_data.iloc[1:][['time']].to_numpy()     # 获得终止时间t1
		
		sub_dt = (sub_t1 - sub_t0) / discrete_t_steps  # 获得迭代时间步长dt
		sub_integrate_t_arr = get_integrate_t_arr(sub_t0, sub_dt, discrete_t_steps)  # 获得积分时间步arr
		
		# 数据拼接
		x0 = vstack_arr(x0, sub_x0)
		x1 = vstack_arr(x1, sub_x1)
		dt = vstack_arr(dt, sub_dt)
		integrate_t_arr = vstack_arr(integrate_t_arr, sub_integrate_t_arr)
	
	# %% 样本和数据转为torch.tensor
	x0 = torch.from_numpy(x0.astype(np.float32))
	x1 = torch.from_numpy(x1.astype(np.float32))
	dt = torch.from_numpy(dt.astype(np.float32))
	integrate_t_arr = torch.from_numpy(integrate_t_arr.astype(np.float32))
	
	dt_tensor = dt.mm(torch.ones(1, cols_n))  # 后续计算会用到

	# %% 搭建网络
	input_size = x0.shape[1] + 1    # 输入为(x, t), 所以shape = dim(x) + dim(t)
	output_size = x0.shape[1]       # 输出为变量x的导数dx/dt, 所以shape = dim(x)
	hidden_sizes = [4 * input_size, 8 * input_size, 8 * output_size, 4 * output_size]

	pd_net = PartialDeriveNet(input_size, hidden_sizes, output_size)

	# %% 指定优化器
	optimizer = torch.optim.Adam(
		pd_net.parameters(),
		lr = lr
	)

	# %% 进行积分
	x_init = x0[0: 1, :]
	t_init = data['time'].min()
	t_final = data['time'].max()
	verity_dt = 0.01

	train_loss_record = []
	plt.figure(figsize = [12, 5])
	pd_net.train()
	for epoch in range(epochs):
		x1_pred, _ = ODESolver(x0, integrate_t_arr, dt_tensor, pd_net)
		train_loss = criterion(x1, x1_pred)

		optimizer.zero_grad()
		train_loss.backward()
		optimizer.step()

		train_loss_record.append(train_loss)

		if (epoch + 1) % 50 == 0:
			print(epoch, train_loss)

		if (epoch + 1) % 200 == 0:
			_, x_records = integrate(x_init, t_init, t_final, verity_dt, pd_net)

			plt.clf()
			plt.subplot(1, 2, 1)
			x1_true_epoch = x1.detach().cpu().numpy()
			x1_pred_epoch = x1_pred.detach().cpu().numpy()
			plt.title('phase portrait')
			plt.scatter(x1_true_epoch[:, 0], x1_true_epoch[:, 1], c = '#1f77b4', s = 6, label = 'true')
			plt.scatter(x1_pred_epoch[:, 0], x1_pred_epoch[:, 1], c = '#ff7f0e', s = 6, label = 'pred')
			plt.legend(loc = 'lower left')

			plt.subplot(1, 2, 2)
			x1_true_epoch = x1.detach().cpu().numpy()
			x1_pred_epoch = x_records.detach().cpu().numpy()
			plt.title('phase portrait')
			plt.scatter(x1_true_epoch[:, 0], x1_true_epoch[:, 1], c = '#1f77b4', s = 6, label = 'true')
			plt.plot(x1_pred_epoch[:, 0, 0], x1_pred_epoch[:, 0, 1], c = '#ff7f0e', linewidth = 1.0, label = 'pred')
			plt.legend(loc = 'lower left')
			plt.pause(1.0)

		# 保存模型
		if (epoch + 1) % 500 == 0:
			print('saving model...')
			pd_net.eval()
			plt.savefig('../../graph/train_effect.png', dpi = 450)
			save_models(pd_net, train_loss_record)
			pd_net.train()

	save_models(pd_net, train_loss_record)
